package utils

import (
	"sync"

	"golang.org/x/time/rate"
)

var FlowLimitManager *FlowLimiter

type FlowLimiter struct {
	FlowLimitMap  map[string]*FlowLimiterItem
	FlowLimitList []*FlowLimiterItem
	Locker        sync.RWMutex
}

type FlowLimiterItem struct {
	ServiceName string
	Limiter     *rate.Limiter
}

func NewFlowLimiter() *FlowLimiter {
	return &FlowLimiter{
		FlowLimitMap:  map[string]*FlowLimiterItem{},
		FlowLimitList: []*FlowLimiterItem{},
		Locker:        sync.RWMutex{},
	}
}

func init() {
	FlowLimitManager = NewFlowLimiter()
}

// 用 rate 包做的限流，根据 ip 限流
func (f *FlowLimiter) GetLimiter(serverName string, qps float64) (*rate.Limiter, error) {
	f.Locker.RLock() // 读锁查找
	if v, ok := f.FlowLimitMap[serverName]; ok {
		f.Locker.RUnlock()
		return v.Limiter, nil
	}
	f.Locker.RUnlock()

	f.Locker.Lock()
	defer f.Locker.Unlock()
	if v, ok := f.FlowLimitMap[serverName]; ok {
		f.Locker.Unlock()
		return v.Limiter, nil
	}

	newLimiter := rate.NewLimiter(rate.Limit(qps), int(qps*3))
	item := &FlowLimiterItem{
		ServiceName: serverName,
		Limiter:     newLimiter,
	}
	f.FlowLimitList = append(f.FlowLimitList, item)
	f.FlowLimitMap[serverName] = item
	return newLimiter, nil
}
